/*
 * COPYRIGHT:	2010 Christopher Smith, licenced under the GNU GPL v2; see COPYING in the top level directory
 * PROJECT:		sdedit
 * FILE:		netapi.cpp
 * PURPOSE:		netapi support
 *				
 */

namespace Netapi
{

template <class T> struct NetapiHeap
{
	typedef NetapiHeap type;

	T *alloc(size_t length)
	{
		T *ret;
		RaiseOnWin32Error(NetApiBufferAllocate(narrowing_cast<DWORD>(length), reinterpret_cast<PVOID *>(&ret)));
		return ret;
	}

	void free(T *x)
	{
		RaiseOnWin32Error(NetApiBufferFree(x));
	}
};

void free_DsRoleBasic(DSROLE_PRIMARY_DOMAIN_INFO_BASIC *x)
{
	DsRoleFreeMemory(x);
}

PWSTR SkipAnyLeadingBackslashes(PWSTR Str)
{
	if(Str[0] == '\\' && Str[1] == '\\')
		return Str + 2;
	else
		return Str;	
}

auto_take<WCHAR> PrefixNameToUNC(PCWSTR Name)
{
	auto_szpod<WCHAR> NamePrefixed;

	if(Name)
	{
		SIZE_T NamePrefixedCt = wcslen(Name) + 4;
		NamePrefixed.realloc_ct(NamePrefixedCt);
		swprintf_s(NamePrefixed, NamePrefixedCt, L"\\\\%s", Name);
	}

	return NamePrefixed.release();
}

auto_take<WCHAR> PrefixNameToUNC(auto_take<WCHAR> Name)
{
	auto_szpod<WCHAR> NameCopy = Name;
	return PrefixNameToUNC(NameCopy);
}

BOOLEAN AreLabelsSupportedRemotely(PCWSTR HostName)
{
	auto_ptr<WKSTA_INFO_100, NetapiHeap> Wki;

	if(RaiseOnWin32Error(NetWkstaGetInfo(c_cast<PWSTR>(HostName), 100, c_cast<PBYTE*>(Wki.assign_indirect())), FALSE) == ERROR_SUCCESS)
		return Wki->wki100_platform_id == PLATFORM_ID_NT && Wki->wki100_ver_major >= 6;
	else
		return FALSE;
}

auto_take<WCHAR> NetapiLookupDcHost(PCWSTR SearchHost, PCWSTR Domain)
{
	OpText Ot2;
	if(SearchHost)
		Ot2.SetText(L"From host", SearchHost);
	OpText Ot(L"Finding a domain controller for domain with NetGet[Any]DCName", Domain);

	auto_szpod<WCHAR> SearchHostPrefixed = PrefixNameToUNC(SearchHost);

	auto_szpod<WCHAR, NetapiHeap> Result;

	ULONG Status = NetGetAnyDCName(SearchHostPrefixed, Domain, reinterpret_cast<PBYTE *>(Result.assign_indirect()));

	if(Status != ERROR_SUCCESS)
	{
		if(Status == ERROR_NO_SUCH_DOMAIN) //NetGetDCName can be less efficent, but works in more cases-- cases that return NO_SUCH_DOMAIN
			RaiseOnWin32Error(NetGetDCName(SearchHostPrefixed, Domain, reinterpret_cast<PBYTE *>(Result.assign_indirect())));
		else
			RaiseWin32Error(Status);
	}

	auto_szpod<WCHAR> Ret = SkipAnyLeadingBackslashes(Result);

	Ot;
	Ot2;

	return Ret.release();
}

auto_take<WCHAR> DsLookupDcHost(PCWSTR SearchHost, PCWSTR Domain)
{
	OpText Ot2;
	if(SearchHost)
		Ot2.SetText(L"From host", SearchHost);
	OpText Ot(L"Finding a domain controller for domain with DsGetDcName", Domain);

	auto_ptr<DOMAIN_CONTROLLER_INFO, NetapiHeap> Dci;

	RaiseOnWin32Error(DsGetDcName(SearchHost, Domain, NULL, NULL, 0, Dci.assign_indirect()));

	auto_szpod<WCHAR> Ret = SkipAnyLeadingBackslashes(Dci->DomainControllerName);

	Ot;
	Ot2;

	return Ret.release();
}

auto_take<WCHAR> LookupDcHost(PCWSTR SearchHost, PCWSTR Domain)
{
	__try
	{
		return DsLookupDcHost(SearchHost, Domain);
	}
	__except(exception_code() == STATUS_VCPP_PROC_NOT_FOUND ? EXCEPTION_EXECUTE_HANDLER : EXCEPTION_CONTINUE_SEARCH)
	{//if the DsGetDcName function isn't available in this version, use the safe Net* equivalent instead
		return NetapiLookupDcHost(SearchHost, Domain);
	}
}

struct AccountEnumerator : Lookup::Enumerator<SID>
{
	ULONG AcctFilter;
	auto_szpod<WCHAR> SystemName;
	auto_szpod<SID> DomainSid;

	AccountEnumerator(PCWSTR SystemName, SID const *DomainSid, ULONG AcctFilter) :
		SystemName(PrefixNameToUNC(SystemName)),
		DomainSid(DomainSid),
		AcctFilter(AcctFilter)
	{}

	AccountEnumerator(Lookup::DOMAIN_INFO const &Di, ULONG AcctFilter) :
		DomainSid(Di.DomainSid),
		SystemName(PrefixNameToUNC(Di.Host)),
		AcctFilter(AcctFilter)
	{
		if(!(Di.Flags & Lookup::DOMAIN_INFO_HAS_HOST))
			SystemName = PrefixNameToUNC(LookupDcHost(Di.SearchHost, Di.DomainName));
	}

	virtual void WalkChildren(Lookup::ChildrenWalkDest<SID> &Dest) override
	{} //an account is a leaf

	virtual void WalkEntries(Lookup::EntriesWalkDest<SID> &Dest) override
	{
		CppCancelOneHandlerExistance CppCohe;

		try
		{
			OpText Ot(L"Enumerating groups", SystemName);
			auto_ptr<GROUP_INFO_2, NetapiHeap> GroupInfo;
			ULONG EntryCt, EntryTotalCt;

			RaiseOnWin32Error(NetGroupEnum(SystemName, 2, c_cast<LPBYTE *>(GroupInfo.assign_indirect()), MAX_PREFERRED_LENGTH, &EntryCt, &EntryTotalCt, NULL));

			for(ULONG i = 0; i < EntryCt; i++)
			{
				auto_szpod<SID> Sid = AddRidToSid(DomainSid, GroupInfo[i].grpi2_group_id);

				Dest.EntryProc(GroupInfo[i].grpi2_name, *Sid);
			}

			Ot;
		}
		catch(CancelOneException)
		{}

		try
		{
			OpText Ot(L"Enumerating users", SystemName);
			auto_ptr<USER_INFO_20, NetapiHeap> UserInfo;
			ULONG EntryCt, EntryTotalCt;

			RaiseOnWin32Error(NetUserEnum(SystemName, 20, AcctFilter, c_cast<LPBYTE *>(UserInfo.assign_indirect()), MAX_PREFERRED_LENGTH, &EntryCt, &EntryTotalCt, NULL));

			for(ULONG i = 0; i < EntryCt; i++)
			{
				auto_szpod<SID> Sid = AddRidToSid(DomainSid, UserInfo[i].usri20_user_id);

				Dest.EntryProc(UserInfo[i].usri20_name, *Sid);
			}
			Ot;
		}
		catch(CancelOneException)
		{}
	}
};

struct DomainAccountEnumerator : Lookup::DomainAccountEnumerator
{
	ULONG AcctFilter;

	DomainAccountEnumerator(ULONG AcctFilter) :
		AcctFilter(AcctFilter)
	{}

	DomainAccountEnumerator(SDE_OPTS Opts) :
		AcctFilter(static_cast<ULONG>(Opts.Opts[OiAcctFilter].Enum->Choice))
	{}


	virtual Lookup::Enumerator<SID> *NewEnum(Lookup::DOMAIN_INFO const &Di) override
	{
		return new AccountEnumerator(Di, AcctFilter);
	}
};

//return is a mask of DOMAIN_INFO_* bits
ULONG GetPrimaryDomainFlagsFromDsRoleInfo(DSROLE_PRIMARY_DOMAIN_INFO_BASIC *DsDomainInfo)
{
	switch(DsDomainInfo->MachineRole)
	{
		case DsRole_RoleStandaloneWorkstation:
		case DsRole_RoleStandaloneServer:
			return Lookup::DOMAIN_INFO_LOCALDOMAIN; //no primary domain

		case DsRole_RoleMemberWorkstation:
		case DsRole_RoleMemberServer:
		case DsRole_RolePrimaryDomainController:
		case DsRole_RoleBackupDomainController:
			return Lookup::DOMAIN_INFO_NETDOMAIN | (DsDomainInfo->Flags & DSROLE_PRIMARY_DS_RUNNING ? Lookup::DOMAIN_INFO_NTDS : Lookup::DOMAIN_INFO_NOT_NTDS);

		default:
			return 0;
	}
}

struct DsDomainEnumerator : Lookup::DomainEnumerator
{
	auto_szpod<WCHAR> SearchHost;
	ULONG DsnameOpts;

	DsDomainEnumerator(SDE_GENERIC_OPTS const &GenOpts, SDE_OPTS Opts)
	{
		SDE_OPT *Opt = Opts.Opts;

		DsnameOpts = narrowing_cast<ULONG>(Opt[OiGetDsName].Enum->Choice);
		
		SearchHost = (Opt[OiHost].Flags & SDE_OPT_SET) ? Opt[OiHost].String : GenOpts.Host;

		if(!SearchHost)
		{
			PCWSTR Domain;

			Domain = (Opt[OiDomain].Flags & SDE_OPT_SET) ? Opt[OiDomain].String : GenOpts.Domain;

			if(Domain)
			{
				OpText Ot(L"Finding a domain controller for domain with DsGetDcName", Domain);

				//we'll have to find a host for the specified domain
				auto_ptr<DOMAIN_CONTROLLER_INFO, NetapiHeap> Dci;
				RaiseOnWin32Error(DsGetDcName(SearchHost, Domain, NULL, NULL, DsnameOpts, Dci.assign_indirect()));

				SearchHost = SkipAnyLeadingBackslashes(Dci->DomainControllerName);

				Ot;
			}//else just leave it as null for local machine
		}
	}

	virtual void WalkEntries(Lookup::EntriesWalkDest<Lookup::DOMAIN_INFO> &Dest) override
	{
		//list search host
		{
			auto_ptr<USER_MODALS_INFO_2, NetapiHeap> NetDomainInfo; //for domain sid
			RaiseOnWin32Error(NetUserModalsGet(SearchHost, 2, c_cast<PBYTE *>(NetDomainInfo.assign_indirect())));

			auto_free<DSROLE_PRIMARY_DOMAIN_INFO_BASIC, &free_DsRoleBasic> DsDomainInfo;
			RaiseOnWin32Error(DsRoleGetPrimaryDomainInformation(SearchHost, DsRolePrimaryDomainInfoBasic, c_cast<PBYTE *>(DsDomainInfo.assign_indirect())));
			PWSTR DomainName = DsDomainInfo->DomainNameDns ? DsDomainInfo->DomainNameDns : DsDomainInfo->DomainNameFlat;

			switch(DsDomainInfo->MachineRole)
			{
			case DsRole_RoleStandaloneWorkstation: //just the local comp
			case DsRole_RoleStandaloneServer:
			case DsRole_RoleMemberWorkstation: //has both local and primary domain
			case DsRole_RoleMemberServer:
				{//report local domain
					Lookup::DOMAIN_INFO Di = {Lookup::DOMAIN_INFO_HAS_HOST | Lookup::DOMAIN_INFO_LOCALDOMAIN,
						c_cast<SID *>(NetDomainInfo->usrmod2_domain_id), NetDomainInfo->usrmod2_domain_name, NetDomainInfo->usrmod2_domain_name};
					Dest.EntryProc(NetDomainInfo->usrmod2_domain_name, Di);
				}
				if(DsDomainInfo->MachineRole == DsRole_RoleStandaloneWorkstation ||
				   DsDomainInfo->MachineRole == DsRole_RoleStandaloneServer)
					return; //no related domains
				else
				{//enum primary domain
					//get domain host
					auto_ptr<DOMAIN_CONTROLLER_INFO, NetapiHeap> Dci;
					RaiseOnWin32Error(DsGetDcName(SearchHost, DomainName, NULL, NULL, DsnameOpts, Dci.assign_indirect()));
					auto_szpod<WCHAR> PdDcName = SkipAnyLeadingBackslashes(Dci->DomainControllerName);

					//get domain SID
					auto_ptr<USER_MODALS_INFO_2, NetapiHeap> NetPdDomainInfo;
					RaiseOnWin32Error(NetUserModalsGet(PdDcName, 2, c_cast<PBYTE *>(NetPdDomainInfo.assign_indirect())));
					
					//report findings
					Lookup::DOMAIN_INFO Di = {Lookup::DOMAIN_INFO_HAS_HOST | Lookup::DOMAIN_INFO_HAS_SEARCHHOST | Lookup::DOMAIN_INFO_NETDOMAIN |
						(DsDomainInfo->Flags & DSROLE_PRIMARY_DS_RUNNING ? Lookup::DOMAIN_INFO_NTDS : Lookup::DOMAIN_INFO_NOT_NTDS),
						c_cast<SID *>(NetPdDomainInfo->usrmod2_domain_id), DomainName, PdDcName, SearchHost, DsDomainInfo->DomainForestName};
					Dest.EntryProc(DomainName, Di);
				}
				break; //continue to trusted domains

			case DsRole_RolePrimaryDomainController: //system is a DC
			case DsRole_RoleBackupDomainController:
				{//report own domain
					Lookup::DOMAIN_INFO Di = {Lookup::DOMAIN_INFO_HAS_HOST | Lookup::DOMAIN_INFO_NETDOMAIN | Lookup::DOMAIN_INFO_HAS_FORESTNAME |
						(DsDomainInfo->Flags & DSROLE_PRIMARY_DS_RUNNING ? Lookup::DOMAIN_INFO_NTDS : Lookup::DOMAIN_INFO_NOT_NTDS),
						c_cast<SID *>(NetDomainInfo->usrmod2_domain_id), DomainName, SearchHost, NULL, DsDomainInfo->DomainForestName};
					Dest.EntryProc(DomainName, Di);
				}
			}
		}

		auto_ptr<DS_DOMAIN_TRUSTS, NetapiHeap> Domains;
		ULONG DomainCount;

		RaiseOnWin32Error(DsEnumerateDomainTrusts(SearchHost, DS_DOMAIN_DIRECT_INBOUND, Domains.assign_indirect(), &DomainCount));

		for(ULONG i = 0; i < DomainCount; i++)
		{
			if(!Domains[i].DomainSid)
				continue; //no domain sid = no relevance

			PWSTR DomainName = Domains[i].DnsDomainName ? Domains[i].DnsDomainName : Domains[i].NetbiosDomainName;

			Lookup::DOMAIN_INFO Di = {Lookup::DOMAIN_INFO_HAS_SEARCHHOST | Lookup::DOMAIN_INFO_NETDOMAIN,
				c_cast<SID *>(Domains[i].DomainSid), DomainName, NULL, SearchHost};

			Dest.EntryProc(DomainName, Di);
		}
	}

	virtual void WalkChildren(Lookup::ChildrenWalkDest<Lookup::DOMAIN_INFO> &Dest) override
	{}
};

struct NetapiModule : Module::ModuleObject
{
	auto_SDE_OPTS ModOpts;

	NetapiModule() :
		ModOpts(StaticNetapiOpts)
	{}

	PCWSTR ModuleName() override
	{
		return L"Netapi";
	}

	virtual SDE_OPTS &ModuleOptions() override
	{
		return ModOpts;
	}

	virtual void OpenResources(Module::OPEN_STATUS &OpenStatus, SDE_GENERIC_OPTS const &GOpts, Module::ModuleResourceDest &Res, UCHAR Privileges, BOOLEAN ObjectWanted) override
	{
		auto_refct<Lookup::DomainEnumerator> DsDomainEnum = new Netapi::DsDomainEnumerator(GOpts, ModOpts);

		Res.DomainEnumerator(L"Netapi", DsDomainEnum);

		auto_refct<Lookup::DomainAccountEnumerator> NetapiDomainAcctEnum = new Netapi::DomainAccountEnumerator(ModOpts);

		Res.DomainSidEnumerator(NetapiDomainAcctEnum);

		if(ObjectWanted)
			OpenStatus |= Module::OBJECT_UNKNOWN_TYPE;
	}
};

}//namespace Netapi